# -*- coding: utf-8 -*-
# @Time    : 2020/6/17 下午11:20
# @Author  : caotian
# @FileName: matplotlibtrain.py
# @Software: PyCharm
# ***
# 使用Matplotlib库绘制损失随训练下降的曲线图
# 将训练的批次编号作为X轴坐标，该批次的训练损失作为Y轴坐标。
#
# 训练开始前，声明两个列表变量存储对应的批次编号(iters=[])和训练损失(losses=[])。
#
# 随着训练的进行，将iter和losses两个列表填满。
#
# 训练结束后，将两份数据以参数形式导入PLT的横纵坐标。
#
# plt.xlabel("iter", fontsize=14)，plt.ylabel("loss", fontsize=14)
# 最后，调用plt.plot()函数即可完成作图。
# ***
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D,Pool2D,Linear
import random
import json
import gzip
import os
import sys
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('TkAgg')
curpath=os.path.abspath(os.curdir)
sys.path.append(curpath)
import optimizationdata as od
import optimizationmodel as om
with fluid.dygraph.guard():
    model =om.MNIST()
    model.train()
    optimizer=fluid.optimizer.SGDOptimizer(learning_rate=0.01,parameter_list=model.parameters())
    epoch_num=10
    iter=0
    iters=[]
    losses=[]
    train_loader=od.load_data('train')
    for epoch_id in range(epoch_num):
        for batch_id,data in enumerate(train_loader()):
            # 准备数据，变得更加简洁
            image_data, label_data = data
            image = fluid.dygraph.to_variable(image_data)
            label = fluid.dygraph.to_variable(label_data)

            # 前向计算的过程，同时拿到模型输出值和分类准确率
            predict, acc = model(image, label)

            # 计算损失，取一个批次样本损失的平均值
            loss = fluid.layers.cross_entropy(predict, label)
            avg_loss = fluid.layers.mean(loss)

            # 每训练了100批次的数据，打印下当前Loss的情况
            if batch_id % 100 == 0:
                print("epoch: {}, batch: {}, loss is: {}, acc is {}".format(epoch_id, batch_id, avg_loss.numpy(),acc.numpy()))
                iters.append(iter)
                losses.append(avg_loss.numpy())
                iter=iter + 100
            # 后向传播，更新参数的过程
            avg_loss.backward()
            optimizer.minimize(avg_loss)
            model.clear_gradients()
    fluid.save_dygraph(model.state_dict(),'mnist-model')

    plt.figure()
    plt.title('train loss',fontsize=24)
    plt.xlabel('iter',fontsize=14)
    plt.ylabel('loss',fontsize=14)
    plt.plot(iters,losses,color='red',label='train loss')
    plt.grid()
    plt.show()